'''
这里分别构建主要的深度学习的问题
'''
import matplotlib.pyplot as plt
import numpy as np

from skimage.data import astronaut
from skimage.color import rgb2gray
from skimage import filters as filters
from skimage.segmentation import felzenszwalb, slic, quickshift, watershed
from skimage.segmentation import mark_boundaries
from skimage.util import img_as_float
import torch
from models.get_models import get_models
import argparse
from utils.utils import get_optimizer,get_scheduler,save_ckpt,save_ckpt_bestmiou,load_ckpt
from utils.loginfomation import loginfomation,WriterAccurary,WriterSummary
from utils.loss import *

from datasetloader.get_datasetLoader import *
import time
from utils.metrics import runningScore
from utils.loss import get_lossfunction
from torchvision.utils import make_grid
import os
import torch.nn.functional as F
import sys
import threading
from time import sleep

from scipy import io




'''
1、直接指出当前的神经网络中存在的问题
2、构建新的算法
3、对比算法结果
'''
def predict(img_path,label_path,datasetName,modelName,cls_num,channels_num=3,redmodelpath=""):
    img=np.load(img_path)
    label=np.load(label_path)
    dataset_path=getRootdirFromDatasetName(datasetName)
    test_transpose=get_augmentations('test',(256,256),dataset_path,cls_num=dataset_path['cls_num'],VI_enable=False)
    
    output_= {'img':img,"seg":label,"label":label,"name":img_path,'source_image':np.copy(img)}
    sample=test_transpose(output_)
    model=get_models(modelName,num_class=cls_num,channels_num=3)
    startepoch,model,best_iou=load_ckpt(model, None, redmodelpath,'cuda')

    model.cuda()
    model.eval()
    with torch.no_grad():
        image,label,source_img,seg=sample['img'],  sample['label'],sample["source_image"],sample['seg']
        image = image.cuda().reshape(1,3,512,512)
        # 
        outputs = model(image.cuda())
        pred_softmax=F.softmax(outputs).cpu()
        # 
        return outputs.cpu().numpy()[0,:,:,:],pred_softmax.cpu().numpy()[0,:,:,:]

def superPixel(img_path):
    img=np.load(img_path)
    segments_fz = felzenszwalb(img, scale=100, sigma=0.5, min_size=50)
    segments_slic = slic(img, n_segments=250, compactness=10, sigma=1)
    segments_quick = quickshift(img, kernel_size=3, max_dist=6, ratio=0.5)
    gradient = filters.sobel(rgb2gray(img))
    segments_watershed = watershed(gradient, markers=250, compactness=0.001)
    return segments_fz,segments_slic,segments_quick,segments_watershed

def problemShow(imgpath,labelpath,fcn_out,fcn_pred,seg_out,seg_pred,unet_out,unet_pred,segments_fz,segments_slic,segments_quick,segments_watershed,superpixel_mask_all):
    img=np.load(imgpath)
    label=np.load(labelpath)*255
    fig, ax = plt.subplots(7, 3, figsize=(10, 10), sharex=True, sharey=True)
    plt.subplots_adjust(left=0, bottom=0, right=0.01, top=0.01, wspace=0, hspace=0)
    # 
    fcn_max=np.argmax(fcn_pred,axis=0)
    seg_max=np.argmax(seg_pred,axis=0)
    unet_max=np.argmax(unet_pred,axis=0)
    # 原图7
    ax[0, 0].imshow(img)
    ax[1, 0].imshow(img)
    ax[2, 0].imshow(img)
    ax[3, 0].imshow(img)
    ax[4, 0].imshow(img)
    ax[5, 0].imshow(img)
    ax[6,0].imshow(img)
    ax[0, 1].imshow(mark_boundaries(img, fcn_max))
    ax[1, 1].imshow(mark_boundaries(img, seg_max))
    ax[2, 1].imshow(mark_boundaries(img, unet_max))
    ax[3, 1].imshow(mark_boundaries(img, segments_slic))
    ax[4, 1].imshow(mark_boundaries(img, segments_quick))
    ax[5, 1].imshow(mark_boundaries(img, segments_watershed))
    ax[6,1].imshow(mark_boundaries(img,superpixel_mask_all))
    ax[0, 2].imshow(label)
    ax[1, 2].imshow(label)
    ax[2, 2].imshow(label)
    ax[3, 2].imshow(label)
    ax[4, 2].imshow(label)
    ax[5, 2].imshow(label)
    ax[6,2].imshow(label)
    for a in ax.ravel():
        a.set_axis_off() 
        plt.subplots_adjust(left=0, bottom=0, right=0.01, top=0.01, wspace=0, hspace=0)
    
    plt.tight_layout()
    plt.show()
    plt.savefig(os.path.join(os.getcwd(),"result.png"),dpi=300)
    print(os.path.join(os.getcwd(),"result.png"))
def arigmation(fcn_pred,superpixel,label_path,imgpath):
    label=np.load(label_path)
    img=np.load(imgpath)
    # 算法流程
    # step1 获得不同的模型计算的pred计算结果
    pred_1=fcn_pred[1,:,:] # 第2类的计算结果
    pred_mask=pred_1>0.7 # 设置阈值
    # 计算梯度
    gradient=np.gradient(pred_1)
    gradientx=gradient[0] # x
    gradienty=gradient[1] # y 
    # 再次求导 确定边界
    dd_x=filters.sobel(gradientx)
    segments_dd_x = watershed(dd_x, markers=250, compactness=0.001)
    dd_y=filters.sobel(gradienty)
    segments_dd_y=watershed(dd_y, markers=250, compactness=0.001)
    io.savemat(os.path.join(os.getcwd(),'gene_features.mat'), {'gradientx': gradientx,'gradienty':gradienty,
                                    'fcn_pred':pred_1,'superpixel':superpixel})
    # step2 超像素
    superpixel_mask=superpixel*pred_mask # superpixel_result
    unique_id=np.unique(superpixel_mask) # 获得所有超像素子图id
    # step2 判断子图的类别归属，这里分为2种，内部与边界
    # 设置阈值对于>0.5 认为就是可能的值
    edge_id=[]
    counter_pixel={}
    superpixel_mask_all=superpixel_mask*0
    for i in unique_id:
        if i==0: # 这是背景
            continue
        # 获取超像素单元
        superpixel_sig=superpixel==i #  单个像元的大小
        pred_sig=pred_1*superpixel_sig # 单个像元的预测值
        superpixel_sig_count=np.sum(np.sum(superpixel_sig)) # 单个超像元的数量
        pred_sig_mean=np.sum(np.sum(pred_sig))/superpixel_sig_count
        # 计算阈值
        if pred_sig_mean>0.5:
            superpixel_mask_all=superpixel_mask_all+superpixel_sig

    superpixel_mask=superpixel_mask>0
    '''
    fig, ax = plt.subplots(2, 3, figsize=(10, 10), sharex=True, sharey=True)
    ax[0, 0].matshow(gradientx)
    ax[0, 1].matshow(gradienty)
    ax[0, 2].matshow(mark_boundaries(pred_1, superpixel))
    ax[1,0].matshow(mark_boundaries(img, segments_dd_x))
    ax[1,1].matshow(mark_boundaries(img, segments_dd_y))
    ax[1,2].matshow(mark_boundaries(img, superpixel_mask_all))
    for a in ax.ravel():
        a.set_axis_off() 
        plt.subplots_adjust(left=0, bottom=0, right=0.01, top=0.01, wspace=0, hspace=0)
    
    plt.tight_layout()
    plt.show()
    '''
    return superpixel_mask_all



def get_edge(mask,pred_1,gradientx_mask,gradienty_mask,dd_x):

    h,w=mask.shape
    for r in range(h):
        for w in range(w):
            # 确定当前的参数信息
            # 
            pass
    pass


# 问题主流程

def mainprocess():
    imgpath='/media/gis/databackup/ayc/modellist/dataset/nanchang/E512/RGB123/test/img/1818.npy'
    labelpath='/media/gis/databackup/ayc/modellist/dataset/nanchang/E512/RGB123/test/label/1818.npy'
    fcn_model_path='/media/gis/databackup/ayc/model_hastrained_log/RGB123_fcn_result/checkpoint/Bestest_mIou_ckpt_fcn_dataName_RGB123.pth'
    seg_model_path='/media/gis/databackup/ayc/model_hastrained_log/RGB123_segnet_result/checkpoint/Bestest_mIou_ckpt_segnet_dataName_RGB123.pth'
    unet_model_path='/media/gis/databackup/ayc/model_hastrained_log/RGB123_unet_result/checkpoint/Bestest_mIou_ckpt_unet_dataName_RGB123.pth'
    # 模型预测阶段
    fcn_out,fcn_pred=predict(imgpath,labelpath,'RGB123','fcn',2,channels_num=3,redmodelpath=fcn_model_path)
    seg_out,seg_pred=predict(imgpath,labelpath,'RGB123','segnet',2,channels_num=3,redmodelpath=seg_model_path)
    unet_out,unet_pred=predict(imgpath,labelpath,'RGB123','unet',2,channels_num=3,redmodelpath=unet_model_path)
    # 超像素结果
    segments_fz,segments_slic,segments_quick,segments_watershed=superPixel(imgpath) # 超像素构建结果
    # 展示结果对比
    superpixel_mask_all=arigmation(unet_pred,segments_quick,labelpath,imgpath)
    problemShow(imgpath,labelpath,fcn_out,fcn_pred,seg_out,seg_pred,unet_out,unet_pred,segments_fz,segments_slic,segments_quick,segments_watershed,superpixel_mask_all)
    pass


if __name__=="__main__":
    mainprocess()


